/**
 * 
 */
package com.hu.ta.assignment6.beans;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

/**
 * @author cage@g.com
 */
public class VSMWordList {
	private Map<String, Token> list;
	private int countSpamDocs;
	private int countHamDocs;

	private static String readString(InputStream is) throws IOException {
		String ret = "";
		int c;
		while ((c = is.read()) != '\0') {
			ret += (char) c;
		}
		return ret;
	}

	public VSMWordList(InputStream is) throws IOException {
		byte[] bytes4 = new byte[4];
		int num;
		is.read(bytes4);
		is.read(bytes4);
		is.read(bytes4);
		num = ByteBuffer.wrap(bytes4).getInt();
		list = new HashMap<String, Token>(num);
		countHamDocs = ByteBuffer.wrap(bytes4).getInt();
		countSpamDocs = ByteBuffer.wrap(bytes4).getInt();
		for (int i = 0; i < num; i++) {
			String key = readString(is);
			is.read(bytes4); // Spam Integer
			is.read(bytes4); // Ham Integer

			list.put(key, new Token(key, ByteBuffer.wrap(bytes4).getInt(),
					ByteBuffer.wrap(bytes4).getInt()));
		}
	}

	public VSMWordList(VSMWordList a) {
		this();
		for (Entry<String, Token> e : a.list.entrySet()) {
			add(e.getKey(), e.getValue().getCount(EmailClass.HAM),
					EmailClass.HAM);
			add(e.getKey(), e.getValue().getCount(EmailClass.SPAM),
					EmailClass.SPAM);
		}
	}

	public void serializeTo(OutputStream os) throws IOException {
		// Anzahl der Wörter
		// wort, abs häufigkeit, rel häufigkeit
		byte[] bytes4 = new byte[4];
		os.write(ByteBuffer.wrap(bytes4).putInt(list.size()).array());
		os.write(ByteBuffer.wrap(bytes4).putInt(countHamDocs).array());
		os.write(ByteBuffer.wrap(bytes4).putInt(countSpamDocs).array());
		for (Entry<String, Token> e : list.entrySet()) {
			os.write(e.getKey().getBytes());
			os.write('\0');
			os.write(ByteBuffer.wrap(bytes4)
					.putInt(e.getValue().getCount(EmailClass.SPAM)).array());
			os.write(ByteBuffer.wrap(bytes4)
					.putInt(e.getValue().getCount(EmailClass.HAM)).array());
		}
	}

	public VSMWordList() {
		list = new HashMap<String, Token>();
		countHamDocs = 0;
		countSpamDocs = 0;
	}

	public void incDoc(EmailClass clas) {
		switch (clas) {
		case HAM:
			countHamDocs++;
			break;
		case SPAM:
			countSpamDocs++;
			break;
		default:
			break;
		}
	}

	public void add(String token, EmailClass clas) {
		add(token, 1, clas);
	}

	public void add(String token, int num, EmailClass clas) {
		token = token.trim().toLowerCase();
		if (null == list.get(token)) {
			switch (clas) {
			case HAM:
				list.put(token, new Token(token, 0, 1));
				return;
			case SPAM:
				list.put(token, new Token(token, 1, 0));
				return;
			default:
				return;
			}
		}
		list.get(token).incNum(clas);
	}

	public boolean contains(String token) {
		token = token.trim().toLowerCase();
		return null != list.get(token);
	}

	public int getCount(String token) {
		return getCount(token, EmailClass.HAM)
				+ getCount(token, EmailClass.SPAM);
	}

	public int getCount(String token, EmailClass clas) {
		Token t = list.get(token.toLowerCase().trim());
		if (t != null) {
			return t.getCount(clas);
		} else {
			return 0;
		}
	}

	public double getProbability(String token, EmailClass clas) {
		Token t = list.get(token.toLowerCase().trim());
		if (t != null) {
			return t.getProbability(clas);
		} else {
			return 0.0;
		}
	}

	public void removeStopWords(List<String> words) {
		for (String w : words) {
			list.remove(w);
		}
	}

	public List<String> getStopWords(int num) {
		List<Tuple<String, Integer>> rankedList = new ArrayList<Tuple<String, Integer>>();
		for (Entry<String, Token> e : list.entrySet()) {
			int idx = 0;
			for (Tuple<String, Integer> t : rankedList) {
				if (t.getY() <= e.getValue().getCount(EmailClass.HAM)) {
					break;
				}
				if (idx >= num) {
					idx = -1;
					break;
				}
				idx++;
			}
			if (idx != -1) {
				rankedList.add(idx, new Tuple<String, Integer>(e.getKey(), e
						.getValue().getCount(EmailClass.HAM)));
			}
		}

		List<String> ret = new ArrayList<String>(rankedList.size());
		int idx = 0;
		for (Tuple<String, Integer> t : rankedList) {
			if (idx >= num) {
				break;
			}
			ret.add(t.getX());
			idx++;
		}

		return ret;
	}

	public double getDocCount(EmailClass clas) {
		switch (clas) {
		case HAM:
			return countHamDocs;
		case SPAM:
			return countSpamDocs;
		default:
			return 0;
		}
	}

	public double getDocProbability(EmailClass clas) {
		switch (clas) {
		case HAM:
			return countHamDocs / (countSpamDocs + countHamDocs);
		case SPAM:
			return countSpamDocs / (countSpamDocs + countHamDocs);
		default:
			return 0;
		}
	}
}
